#!/usr/bin/env python

"""mysql-replace-tables-live
Replace live, in use, MySQL tables without restarting or interrupting
dependent applications.
http://code.google.com/p/mysql-replace-tables-live/
By Richard Bronosky, offered under the MIT License
http://www.opensource.org/licenses/mit-license.php

Sometimes you find yourself needing to replace tables on production servers.
While the ideal solution is to take the server out of rotation and replace the
table properly, that's not always feasible. By default, the app replaces the
table layout (.frm file), the data (.MYD file), and indexes (.MYI file). But
the simplicity of Python allows you to customize it to meet your needs. MySQL
is all about freedom and flexibility. This little app lets you make the most
of it.
"""

from os import path
import sys
import shutil
import re

def main():
    """This main routine outlines the steps to be taken when this script is
    executed directly.  No functions are called if this script is imported
    into another.
    """

    (options, files) = parseCall()
    (files, tables) = validateTableFiles(files)
    mysql=MySQL(options.database, options.user, options.password)
    if not options.quiet:
        verifyIntent(options, files, tables, mysql)
    mysql.lock(tables)
    replaceTables(mysql.datadir+mysql.db, files)
    mysql.unlock(tables)
    
def parseCall():
    """Interpret and handle the command line arguments."""

    from optparse import OptionParser
    parser = OptionParser()
    usage = """
%prog [options] file.frm file.MYD file.MYI [ file.frm file.MYD file.MYI ] ...
       (file triplets can be passed in any order)
""".strip()
    parser = OptionParser(usage=usage)
    parser.add_option("-d", "--database",
            help="MySQL database that contains the tables to be replaced")
    parser.add_option("-u", "--user",
            help="MySQL connection username")
    parser.add_option("-p", "--password",
            default="",
            help="MySQL connection password")
    parser.add_option("-q", "--quiet", action="store_true",
            default=False,
            help="Bypass the normal confirmation step and just do the deed.")
    (options, args) = parser.parse_args()
    options.prog = "mysql-replace-tables-live"
    if len(args)==0:
        parser.print_help()
        sys.exit()
    return (options, args)

def validateTableFiles(files, quiet=True):
    """Verifies filenames, not contents."""

    extensions = ['.frm','.MYD','.MYI']
    tables = []
    files.sort()
    """
    This step has a mild inherent weakness in that all the source files passed
    in for a single table must be in the same directory.  To overcome this you
    could sort using a callback that sorts on filename only, and ignores path.
    I felt that was unlikely to be needed and a waste of time.  YMMV
    """
    files_len=len(files)
    # Check number of files
    if files_len<3 or files_len%3!=0:
        TableFilesError('Wrong number of files passed.')

    # Check names of files
    for triplet in range(len(files)/3):
        sieveOfExtensions = extensions[:]
        if not quiet: print "processing table: ",
        for e in range(len(extensions)):
            file = files[triplet*3+e]
            file = path.split(file)[1]
            basename, ext = path.splitext(file)
            # Identify the name of the table being processed.
            if e==0:
                if not quiet: print "%s...\n    " % basename,
                tables.append(basename)
            if not quiet: print "%s " % file,

            # Check basenames
            try:
                if basename!=basename_old:
                    msg = 'The file %s does not fit into the %s collection of table files.'
                    TableFilesError(msg % (file, basename_old))
            except UnboundLocalError:
                pass
            basename_old=basename

            # Check extensions
            try:
                del sieveOfExtensions[sieveOfExtensions.index(ext)]
            except ValueError:
                msg = 'The file %s does not have a valid extension (%s).'
                TableFilesError(msg % (file, '|'.join(extensions)))

        del basename_old
        if not quiet: print ''
    return (files, tables)

class MySQL(object):
    """This is a wrapper that trivializes all db related tasks."""

    def __init__(self, db, user, password='', host='localhost'):
        """Connect, create cursor, define SQL, and query for datadir"""

        # future versions will allow you to modify remote hosts with ssh/scp
        import MySQLdb
        self.host = host
        self.db = db
        self.user = user
        self.con = MySQLdb.connect(
                user = user,
                passwd = password,
                host = host,
                db = db)
        self.cur = self.con.cursor()
        self.sql_flush = "FLUSH TABLES"
        self.sql_unlock = "UNLOCK TABLES"
        self.sql_lock = "LOCK TABLES %s WRITE;"
        self.sql_lock_delimiter = " WRITE, "
        self.datadir = self.get_datadir()
    def get_datadir(self):
        """Query the server for its datadir variable."""

        self.cur.execute("show variables where Variable_name='datadir'")
        return self.cur.fetchall()[0][1]
    def lock(self, tables, flush=1):
        """Write lock the tables passed an optionally flush tables"""

        sql = self.sql_lock % self.sql_lock_delimiter.join(tables)
        print sql
        self.cur.execute(sql)
        if flush:
            sql = self.sql_flush
            print sql
            self.cur.execute(sql)
    def unlock(self, tables, flush=1):
        """Unlock tables"""

        if flush:
            sql = self.sql_flush
            print sql
            self.cur.execute(sql)
        sql = self.sql_unlock
        print sql
        self.cur.execute(sql)

def verifyIntent(options, files, tables, mysql):
    """Describes in detail the actions to be taken, and gives the user a
    chance to bail out.
    """

    msg = """
%s
  1. The MySQL connection has been made using username: %s
  2. The MyISAM table files were passed for: %s
  3. The table%s above will be WRITE LOCKed.
  4. FLUSH TABLES will be called to flush any caches writes to disk.
  5. The table files will be replaced with the files you passed, which are:
       %s
  6. The schema path was deduced by querying the MySQL server, it is:
       %s
  7. FLUSH TABLES will be called again, because it also includes RESET CACHE.
  8. The lock will then be released, and the connection closed.
""".strip()
    plural = (len(tables)>1) and 's' or ''
    print msg % (options.prog, options.user, ', '.join(tables), plural, "\n       ".join(files), mysql.datadir+mysql.db)
    print ""
    if not ask_yes_no('Do wish to proceed (y/[n])?', 'n'):
        sys.exit()
    print ""

def ask_yes_no(prompt,default=None):
    """Asks a question and returns an integer 1/0 (y/n) answer.

    If default is given (one of 'y','n'), it is used if the user input is
    empty. Otherwise the question is repeated until an answer is given.

    An EOF is treated as the default answer.  If there is no default, an
    exception is raised to prevent infinite loops.

    Valid answers are: y/yes/n/no (match is not case sensitive)."""

    answers = {'y':True,'n':False,'yes':True,'no':False}
    ans = None
    while ans not in answers.keys():
        try:
            ans = raw_input(prompt+' ').lower()
            if not ans:  # response was an empty string
                ans = default
        except KeyboardInterrupt:
            pass
        except EOFError:
            if default in answers.keys():
                ans = default
                print
            else:
                raise
            
    return answers[ans]

def replaceTables(db_path, new_files):
    """Copy the files passed into the dir passed."""

    for file in new_files:
        print "       %s -> %s" % (file, path.join(db_path,path.split(file)[1]))
        shutil.copy(file, db_path)

def TableFilesError(msg):
    """A general handler for errors with the files passed in."""

    print "TableFilesError: %s" % msg
    sys.exit()

if __name__=='__main__':
    main()
